import os, sys

import OpenEXR
import Imath
import math
import time

import numpy as np
from numpy import array
import myextension
from readEXR import *
from writeEXR import *
from MRF_Utils import *


# this function will return the indices of the new localImage, whose area is 1.5 times bigger than the bounding box of the user chosen area, so that in main program on can do indexing quickly
def localImage(labels, Nth):

    tmp = np.where(labels==Nth)

    len_X = np.amax(tmp[0]) - np.amin(tmp[0])
    
    max_X = np.amax(tmp[0]) + len_X*1.5
    min_X = np.amin(tmp[0]) - len_X*1.5

    
    len_Y = np.amax(tmp[1]) - np.amin(tmp[1])
    
    max_Y = np.amax(tmp[1]) + len_Y*1.5
    min_Y = np.amin(tmp[1]) - len_Y*1.5

    return min_X,max_X,min_Y,max_Y

if __name__ == "__main__":
    
    if len(sys.argv)  < 2:
        print "no image input"
        sys.exit(1)
    
    image = sys.argv[1]
    n_labels = 2     
    R,G,B,L,size = readExr(image)
    ########
    R =  np.array(R,dtype = np.double)
    G =  np.array(G,dtype = np.double)
    B =  np.array(B,dtype = np.double)
    L =  np.array(L,dtype = np.double)
    ########
    
    print image,size
    #initialisation of labels
    #labels = np.array(np.random.randint(n_labels,size=size),dtype=np.double)
    
    
    labels = np.ones(size,dtype=np.double)
    # sunflower
    #labels[115:293,(492-378):(492-327)] = 0
    #labels[156:264,(492-182):(492-128)] = 0
    #labels[116:303,(492-312):(492-190)] = 0
    
    #eye.exr
    #labels[81:142,(185-103):(185-49)] = 0
    
    #eye_small.exr
    #labels[15:29,(36-20):(36-9)] = 0
    
    #Pixar05.exr
    #labels[119:205,(702-227):(702-63)] = 0
    #labels[84:241,(702-139):(702-122)] = 0
    
    #pixar.exr
    #labels[50:91,(146-92):(146-44)] = 0
    
    #pixar_creation.exr
    #labels[552:615,(511-229):(511-190)] = 0
    
    
    
    #vue1_samll.exr
    labels[1315:1432,(5616-2537):(5616-2317)] = 0
    
    writeEXR("../../images/label0.exr",np.array(labels,dtype=np.float32).T,np.array(labels,dtype=np.float32).T,np.array(labels,dtype=np.float32).T, size)
    
    min_X,max_X,min_Y,max_Y  = localImage(labels,0)
    

    localR = np.array(R[min_X:max_X+1, min_Y:max_Y+1],dtype = np.double)
    localG = np.array(G[min_X:max_X+1, min_Y:max_Y+1],dtype = np.double)
    localB = np.array(B[min_X:max_X+1, min_Y:max_Y+1],dtype = np.double)
    localLabels = np.array(labels[min_X:max_X+1, min_Y:max_Y+1],dtype = np.double)
     
    """
    localR = R
    localB = B
    localG = G
    localLabels = labels
    """
    
    print localR.shape, localLabels.shape
    print localR[0][0],localG[0][0],localB[0][0]
    maxflow = np.finfo(np.float64).max

    
    writeEXR("../../images/label0_local.exr",np.array(localLabels,dtype=np.float32).T,np.array(localLabels,dtype=np.float32).T,np.array(localLabels,dtype=np.float32).T, localLabels.shape)
    
    writeEXR("../../images/localRGB.exr",np.array(localR,dtype=np.float32).T,np.array(localG,dtype=np.float32).T,np.array(localB,dtype=np.float32).T, localR.shape)
    
    for k in xrange(3): 
        
        inversedCovarianceMatrixArray = []
        miuArray = []
        lnCovarMatDet = []
        covarMatrixArray = []
        
        for i in xrange(n_labels):
            
            covarMatrix, x, y, r, g, b = featuresRGB(localR,localG,localB,localLabels,i)
            
            inversedCovarianceMatrixArray.append(np.linalg.inv(covarMatrix))
            
            miuArray.append((x,y,r,g,b))  
            
            lnCovarMatDet.append(np.log(np.sqrt( 32* np.pi* np.pi* np.pi* np.pi* np.pi  * np.linalg.det(covarMatrix))))
            
            covarMatrixArray.append(covarMatrix)
        
        inversedCovarianceMatrixArray   = np.array(inversedCovarianceMatrixArray,dtype = np.double).reshape((n_labels,5,5))                
        miuArray                        = np.array(miuArray,dtype = np.double).reshape((n_labels,5))
        lnCovarMatDet                   = np.array(lnCovarMatDet,dtype = np.double).reshape(n_labels)
        
        flow = myextension.quickGraphCut(n_labels, localR,localG,localB, localLabels, miuArray, inversedCovarianceMatrixArray,lnCovarMatDet)
        if flow < maxflow:
            maxflow = flow
        else:
            pass
            #sys.exit()
    
        labels[min_X:max_X+1, min_Y:max_Y+1] = localLabels
        #labels = localLabels
        writeEXR("../../images/label"+str(k+1)+".exr",np.array(labels,dtype=np.float32).T,np.array(labels,dtype=np.float32).T,np.array(labels,dtype=np.float32).T, labels.shape)
    




